import torch
import numpy as np
import os
from utils.reranking import re_ranking


def euclidean_distance(qf, gf):
    m = qf.shape[0]
    n = gf.shape[0]
    dist_mat = torch.pow(qf, 2).sum(dim=1, keepdim=True).expand(m, n) + \
               torch.pow(gf, 2).sum(dim=1, keepdim=True).expand(n, m).t()
    dist_mat.addmm_(1, -2, qf, gf.t())
    return dist_mat.cpu().numpy()

def cosine_similarity(qf, gf):
    epsilon = 0.00001
    dist_mat = qf.mm(gf.t())
    qf_norm = torch.norm(qf, p=2, dim=1, keepdim=True)  # mx1
    gf_norm = torch.norm(gf, p=2, dim=1, keepdim=True)  # nx1
    qg_normdot = qf_norm.mm(gf_norm.t())

    dist_mat = dist_mat.mul(1 / qg_normdot).cpu().numpy()
    dist_mat = np.clip(dist_mat, -1 + epsilon, 1 - epsilon)
    dist_mat = np.arccos(dist_mat)
    return dist_mat


def eval_func(distmat, q_pids, g_pids, q_camids, g_camids, max_rank=50):
    """Evaluation with market1501 metric
        Key: for each query identity, its gallery images from the same camera view are discarded.
        """
    #g_pids = [101, 102, 103, 104, 105]  # Gallery ID
    num_q, num_g = distmat.shape
    if num_g < max_rank:
        max_rank = num_g
        print("Note: number of gallery samples is quite small, got {}".format(num_g))
    #indices = np.argsort(distmat, axis=1) 返回的是整数索引排序，表示 distmat 每一行（每个查询样本）按从小到大的顺序排序后，原始索引的位置。（从小到大，距离越小越相似）
    indices = np.argsort(distmat, axis=1)
    #indices = np.array([
        #[0, 2, 1, 3],  # Query 1 的排序
        #[1, 2, 3, 0],  # Query 2 的排序
        #[3, 0, 1, 2]   # Query 3 的排序
        #])
    g_pids_sorted = g_pids[indices]
    #print(g_pids_sorted)输入例如下
    # [[101, 103, 102, 104]  # Query 1
    #  [102, 103, 104, 101]  # Query 2
    #  [104, 101, 102, 103]] # Query 3
    #matches 是一个二值矩阵，表示匹配情况：1 表示匹配正确（同 ID）。0 表示匹配错误。
    #: 表示选取所有行（不改变内容）。np.newaxis 增加一个新的维度，使其从 1D 变成 2D。从[101, 102, 103]变成[[101][102][103]]
    matches = (g_pids_sorted == q_pids[:, np.newaxis]).astype(np.int32)
    #.astype(np.int32)之前是一个布尔二维数组，.astype(np.int32)转化为010101
    #print(matches)
    # [[1, 0, 0, 0]
    #  [0, 1, 0, 0]
    #  [0, 0, 1, 0]]

    all_cmc = []#all_cmc 存储所有 Query 的 CMC 曲线，用于最终计算 Rank-1, Rank-5, Rank-10 的准确率。
    all_AP = []#all_AP 存储所有 Query 的 AP（Average Precision，平均精度），用于计算 mAP（mean Average Precision）。
    num_valid_q = 0.  #num_valid_q 统计有多少个有效的 Query 样本，如果某个 Query 在 Gallery 里没有匹配对象，就不算有效。


    for q_idx in range(num_q):
        # get query pid and camid
        q_pid = q_pids[q_idx]
        q_camid = q_camids[q_idx]

        # remove gallery samples that have the same pid and camid with query
        order = indices[q_idx]  # select one row
        #移除画廊集中同一人同一摄像头的，gpids是个一维数组，g_pids[order]排序之后还是一个一维数组，== q_pid之后变成一维布尔数组，另一个摄像机同理，最后两个布尔数组按位与
        remove = (g_pids[order] == q_pid) & (g_camids[order] == q_camid)
        keep = np.invert(remove)

        # compute cmc curve
        # binary vector, positions with value 1 are correct matches
        orig_cmc = matches[q_idx][keep]
        if not np.any(orig_cmc):
            # 如果 orig_cmc 里全是 0，说明 Query 的身份在 Gallery 中没有出现，这个 Query 无效，直接跳过。
            continue
        #cumsum的作用是前n项和
        #arr = np.array([1, 2, 3, 4, 5])
        #print(np.cumsum(arr))
        #[ 1  3  6 10 15]
        cmc = orig_cmc.cumsum()
        cmc[cmc > 1] = 1#确保 CMC 结果是 {0,1} 二值化

        all_cmc.append(cmc[:max_rank])# 只取前 max_rank 结果
        num_valid_q += 1.# 统计有效 Query 数量

        # compute average precision
        # reference: https://en.wikipedia.org/wiki/Evaluation_measures_(information_retrieval)#Average_precision
        num_rel = orig_cmc.sum()# 计算该 Query 在 Gallery 里前max_rank里有多少个正确匹配
        tmp_cmc = orig_cmc.cumsum()
        #tmp_cmc = [x / (i + 1.) for i, x in enumerate(tmp_cmc)]
        y = np.arange(1, tmp_cmc.shape[0] + 1) * 1.0
        tmp_cmc = tmp_cmc / y
        tmp_cmc = np.asarray(tmp_cmc) * orig_cmc#tmp_cmc 现在是 每个排名的 Precision ,orig_cmc 是 二进制匹配标签(tmp_cmc * orig_cmc) 作用：只有 正确匹配（1）的 Precision 会被保留 ,错误匹配（0）的 Precision 变成 0
        AP = tmp_cmc.sum() / num_rel
        all_AP.append(AP)

    assert num_valid_q > 0, "Error: all query identities do not appear in gallery"

    all_cmc = np.asarray(all_cmc).astype(np.float32)#CMC曲线x轴：前n
    all_cmc = all_cmc.sum(0) / num_valid_q
    mAP = np.mean(all_AP)

    return all_cmc, mAP


class R1_mAP_eval():
    def __init__(self, num_query, max_rank=50, feat_norm=True, reranking=False):
        super(R1_mAP_eval, self).__init__()
        self.num_query = num_query
        self.max_rank = max_rank
        self.feat_norm = feat_norm
        self.reranking = reranking

    def reset(self):
        self.feats = []
        self.pids = []
        self.camids = []

    def update(self, output):  # called once for each batch
        feat, pid, camid = output
        self.feats.append(feat.cpu())
        self.pids.extend(np.asarray(pid))
        self.camids.extend(np.asarray(camid))

    def compute(self):  # called after each epoch
        #self.feats 里存储了 所有 batch 的特征。
        #torch.cat() 把所有特征拼接成 一个完整的特征矩阵。
        feats = torch.cat(self.feats, dim=0)
        #normalize() 进行 L2 归一化，确保所有特征向量的模长为 1。
        #归一化可以提高模型的检索效果。
        if self.feat_norm:
            print("The test feature is normalized")
            feats = torch.nn.functional.normalize(feats, dim=1, p=2)  # along channel
        # query
        qf = feats[:self.num_query]#qf (Query 特征): 前 num_query 个样本是 Query。
        q_pids = np.asarray(self.pids[:self.num_query])
        q_camids = np.asarray(self.camids[:self.num_query])
        # gallery
        gf = feats[self.num_query:]#gf (Gallery 特征): 剩下的样本是 Gallery。
        g_pids = np.asarray(self.pids[self.num_query:])

        g_camids = np.asarray(self.camids[self.num_query:])
        if self.reranking:
            print('=> Enter reranking')
            # distmat = re_ranking(qf, gf, k1=20, k2=6, lambda_value=0.3)
            distmat = re_ranking(qf, gf, k1=50, k2=15, lambda_value=0.3)

        else:
            print('=> Computing DistMat with euclidean_distance')
            distmat = euclidean_distance(qf, gf)
        cmc, mAP = eval_func(distmat, q_pids, g_pids, q_camids, g_camids)

        return cmc, mAP, distmat, self.pids, self.camids, qf, gf



